from typing import List, Dict

from PySide6.QtCore import QRect, QPoint
from PySide6.QtGui import QPainter
from PySide6.QtWidgets import QWidget, QPushButton

from gui.reference import SIZE_WIDTH, SIZE_HEIGHT

BTN_WIDTH = 100
BTN_HEIGHT = 30


class TreeNode:
    __slots__ = ['name', 'left', 'top', 'parents', 'children', 'rect', 'btn']

    def __init__(self, name, left, top):
        self.name = name

        self.left = left
        self.top = top

        self.parents: List[TreeNode] = []
        self.children = []

        self.rect = QRect(left, top, BTN_WIDTH, BTN_HEIGHT)

        self.btn: QPushButton | None = None


def tree_to_dict(tree: TreeNode) -> Dict:
    result = dict()
    q = [tree]
    while q:
        node = q.pop(0)
        result[node.name] = node
        q.extend(node.children)

    return result


class TreeMapWidget(QWidget):
    def __init__(self, title, tree_root):
        super().__init__()

        self.setGeometry(300, 300, SIZE_WIDTH, SIZE_HEIGHT)
        self.setFixedSize(SIZE_WIDTH, SIZE_HEIGHT)
        self.setWindowTitle(title)
        self.root: TreeNode = tree_root
        self.source_y = 0

        self.node_map: Dict[str, TreeNode] = tree_to_dict(tree_root)
        for node_name, node in self.node_map.items():
            btn = QPushButton(node.name, self)
            btn.setGeometry(node.rect.x(), node.rect.y() - self.source_y, BTN_WIDTH, BTN_HEIGHT)
            node.btn = btn

        self.display_tree()

        self.btn_reset = QPushButton('回到顶部', self)
        self.btn_reset.setGeometry(SIZE_WIDTH - BTN_WIDTH, 0, BTN_WIDTH, BTN_HEIGHT)
        self.btn_reset.clicked.connect(self.reset_source_y)

    def reset_source_y(self):
        self.source_y = 0
        self.display_tree()

    def display_tree(self):
        for node in self.node_map.values():
            if node is not None:
                node.btn.move(node.rect.x(), node.rect.y() - self.source_y)

        self.update()

    def wheelEvent(self, event):
        super().wheelEvent(event)
        dy = event.angleDelta().y()

        self.source_y -= dy
        if self.source_y < 0:
            self.source_y = 0

        self.display_tree()

    def paintEvent(self, event):
        super().paintEvent(event)
        painter = QPainter(self)
        for node in self.node_map.values():
            if node is not None:
                for parent in node.parents:
                    p1 = QPoint(parent.rect.x() + BTN_WIDTH // 2, parent.rect.y() + BTN_HEIGHT - self.source_y)
                    p2 = QPoint(node.rect.x() + BTN_WIDTH // 2, node.rect.y() - self.source_y)
                    painter.drawLine(p1, p2)
